import json
import asyncio
from typing import Dict, Any, Callable, List

from agent_tools.all_tools import TOOLS_SET
from my_utils.llm_client import llm_chat, llm_chat_parallel
from my_utils.log_recording import Logging

class MessageManager(Logging):
    """ 用于管理用户、助手、系统和工具的消息记录。 """

    def __init__(self, messages: List[Dict[str, Any]] = None):
        super().__init__()
        self.messages = messages if messages is not None else []

    def add_assistant_message(self, message: str):
        self.messages.append({"role": "assistant", "content": message})
        self.log_debug(f"添加了assistant消息")

    def add_user_message(self, message: str):
        self.messages.append({"role": "user", "content": message})
        self.log_debug(f"添加了user消息")

    def add_system_message(self, message: str):
        self.messages.append({"role": "system", "content": message})
        self.log_debug(f"添加了system消息")

    def add_tool_message(self, tool_call_id: str, tool_name: str, content: str):
        tool_message = {
            "role": "tool",
            "tool_call_id": tool_call_id,
            "name": tool_name,
            "content": content
        }
        self.messages.append(tool_message)
        self.log_debug(f"添加了tool消息")

    def add_message(self, message: Dict[str, Any]):
        if "role" not in message or "content" not in message:
            self.log_error("消息缺少'role'或'content'键。")
            raise ValueError("消息必须包含'role'和'content'键。")
        self.messages.append(message)
        self.log_debug(f"添加了{message['role']}消息")


class ToolManager(Logging):
    """工具管理类，用于管理可用的工具及其对应的函数。"""

    def __init__(self):
        super().__init__()
        self.tools: List[Dict[str, Any]] = []
        self.functions: Dict[str, Callable] = {}

    def add_tool(self, tool_schema: Dict[str, Any], func: Callable):
        """ 添加工具 """

        function_info = tool_schema.get("function")
        if not function_info or "name" not in function_info:
            self.log_error("工具架构信息缺少 'function' 或 'name' 键。")
            raise ValueError("工具架构信息必须包含 'function' 和 'name' 键。")

        tool_name = function_info["name"]
        if tool_name in self.functions:
            self.log_warning(f"'{tool_name}'已存在，将被覆盖。")

        self.tools.append(tool_schema)
        self.functions[tool_name] = func
        self.log_info(f"添加了工具: {tool_name}")

    def list_tools(self) -> List[str]:
        """ 列出已有工具 """

        tool_names = [tool["function"]["name"] for tool in self.tools]
        self.log_debug(f"当前可用工具: {tool_names}")
        return tool_names

    def call_tool(self, name: str, **kwargs) -> Any:
        """ 调用指定名称的函数 """

        if name not in self.functions:
            self.log_error(f"没有找到对应的工具函数：{name}")
            raise ValueError(f"没有找到对应的工具函数：{name}")
        try:
            result = self.functions[name](**kwargs)
            self.log_info(f"工具'{name}'的返回结果为：\n{result}")
            return result
        except Exception as e:
            self.log_error(f"调用工具'{name}'时出错: {e}")
            raise


class BaseAgent(Logging):
    """ Agent基础类 """

    def __init__(self, max_epochs: int = 50):
        super().__init__()

        self.tasks = []
        self.total_tokens = 0
        # self.current_epoch = 0 # 移动到 task 内部
        self.max_epochs = max_epochs
        self.tool_manager = ToolManager()

    def _initialize_tools(self, tool_name_list: List[str]):
        """ 初始化工具集 """

        for tool_name in tool_name_list:
            tool_info = TOOLS_SET.get(tool_name)
            if tool_info:
                self.tool_manager.add_tool(tool_info["tool_schema"], tool_info["tool_def"])
            else:
                self.log_warning(f"添加失败！工具'{tool_name}'不存在于TOOLS_SET中，请检查工具名称是否正确")

    def get_last_assistant_contents(self, messages):
        """ 获取助手最后回复的消息内容 """
        result = []
        if not isinstance(messages, list): # 添加检查
             self.log_warning("get_last_assistant_contents 收到非列表类型的 messages")
             return ""
        for message in reversed(messages):
            if isinstance(message, dict) and message.get("role") == "assistant": # 确保 message 是字典
                # 优先获取 content，如果 content 为空或不存在，检查是否有 tool_calls
                content = message.get("content", "")
                tool_calls_str = ""
                if not content and message.get("tool_calls"):
                     calls = message.get("tool_calls", [])
                     if isinstance(calls, list):
                          tool_calls_str_parts = []
                          for tc in calls:
                              if isinstance(tc, dict) and isinstance(tc.get('function'), dict):
                                   fname = tc['function'].get('name', 'unknown_function')
                                   fargs = tc['function'].get('arguments', '{}')
                                   tool_calls_str_parts.append(f"[{fname}:{fargs}]")
                          tool_calls_str = " ".join(tool_calls_str_parts)

                # 如果 content 和 tool_calls_str 都为空，则此 assistant 消息无效，继续往前找
                if content or tool_calls_str:
                     result.insert(0, content or tool_calls_str) # 取其中不为空的那个
                # else: # 如果都为空，可以忽略这条 assistant 消息，继续往前找
                #     pass

            elif isinstance(message, dict) and message.get("role") != "assistant":
                 break
            # 如果 message 不是字典，也停止
            elif not isinstance(message, dict):
                 break

        return "".join(result)


    # 将 run_loop 改为异步函数
    async def run_loop(self, model_name: str = "GLM_4_FLASH", user_inputs: List[str] = None, system_prompt: str = "", temperature: float = 0.01, tool_choice='auto', tool_name_list: List[str] = [], max_wait_times=600, **args):
        """调用Agent，支持单一任务或多任务并行处理 (现在是异步的)"""
        self.tasks = []
        if user_inputs is None:
            user_inputs = []

        if not user_inputs:
            self.log_error("没有输入用户指令！！！")
            raise ValueError("没有输入用户指令！！！")

        # 初始化工具
        self._initialize_tools(tool_name_list)

        # 为每个用户输入创建独立的任务
        for idx, user_input in enumerate(user_inputs):
            memory = MessageManager()
            if system_prompt:
                memory.add_system_message(system_prompt)
            memory.add_user_message(user_input)
            task = {
                "task_id": idx,
                "memory": memory,
                "current_epoch": 0,
                "latest_state": "start" # 初始状态
            }
            self.tasks.append(task)

        self.log_info("Agent开始工作...")

        loop_count = 0 # 新增循环计数器，防止无限循环
        max_loops = self.max_epochs + 5 # 设定一个最大循环次数，略大于 max_epochs

        while loop_count < max_loops:
            loop_count += 1
            active_tasks = [task for task in self.tasks if not self.is_hit_stop(task)]
            if not active_tasks:
                self.log_info("所有任务均已达到停止条件。")
                break  # 所有任务都已停止

            self.log_info(f"Agent 第 {loop_count} 轮循环，处理 {len(active_tasks)} 个活动任务...")

            # 准备需要处理的任务消息
            messages_list = [task['memory'].messages for task in active_tasks]

            # 使用 await 调用异步的 work 方法
            await self.work(active_tasks=active_tasks, messages_list=messages_list, model_name=model_name, temperature=temperature, tool_choice=tool_choice, max_wait_times=max_wait_times, **args)

            # 添加短暂休眠，避免CPU空转
            await asyncio.sleep(0.1)

        if loop_count >= max_loops:
             self.log_warning(f"Agent 循环达到最大次数 {max_loops}，强制终止。")

        self.log_info("Agent已完成所有任务的工作。")

        # 返回结果的逻辑保持不变，但 get_last_assistant_contents 可能需要微调以处理 tool_calls
        return [
            self.get_last_assistant_contents(task.get('memory', {}).messages)
            for task in self.tasks if isinstance(task, dict)
        ]


    async def work(self, active_tasks: list, messages_list: list, model_name: str, temperature: float, tool_choice, max_wait_times: int = 600,  **args):
        """ 调用llm回答"""
        if not active_tasks: # 添加检查
            self.log_info("没有活动任务需要处理。")
            return

        # 判断是单一任务还是多任务并行
        if len(active_tasks) == 1:
            # 单一任务处理
            task = active_tasks[0]
            self.log_info(f"任务 {task['task_id']} 开始执行第 {task['current_epoch'] + 1} 轮工作 (单任务模式)。")
            try:
                # 如果 llm_chat 也是异步的，需要 await asyncio.to_thread(llm_chat, ...) 或直接 await llm_chat(...)
                response = llm_chat(
                    messages=messages_list[0],
                    model_name=model_name,
                    temperature=temperature,
                    ** ({"tools": self.tool_manager.tools} if self.tool_manager.tools else {}),
                    ** ({"tool_choice": tool_choice} if self.tool_manager.tools else {}),
                    **args
                )

                # --- 处理响应 ---
                if not response or not hasattr(response, 'choices') or not response.choices:
                     self.log_error(f"任务 {task['task_id']} 调用 llm_chat 返回无效响应: {response}")
                     task['latest_state'] = "stop" # 设为停止状态
                     task['current_epoch'] += 1 # 增加轮次计数避免死循环
                     return # 提前返回，不处理无效响应

                choice = response.choices[0]
                task['latest_state'] = choice.finish_reason # 更新状态
                result = choice.message
                result_dump = result.model_dump()

                # 更新内存 - 找到对应的 task 并更新 (这部分逻辑可以优化，直接用 task 变量)
                # task['memory'].add_message(result_dump) # 直接更新当前 task 的 memory

                # 更新 self.tasks 中的对应项 (如果 task 是副本的话需要这样做，但 active_tasks 里的应该是引用)
                for t in self.tasks:
                     if t['task_id'] == task['task_id']:
                          t['memory'].add_message(result_dump)
                          t['latest_state'] = task['latest_state']
                          t['current_epoch'] = task['current_epoch'] # 确保 epoch 也同步（虽然下面会加1）
                          break

                # 处理工具调用
                if hasattr(result, 'tool_calls') and result.tool_calls:
                    self._handle_tool_calls(result.tool_calls, task)

                # 记录 token
                if hasattr(response, 'usage') and hasattr(response.usage, 'total_tokens'):
                    # if 'total_tokens' in response.usage: # 更安全的检查
                    self.total_tokens += response.usage.total_tokens
                    self.log_info(f"任务 {task['task_id']} 本次回复 token: {response.usage.total_tokens}, 总 token: {self.total_tokens}")

                task['current_epoch'] += 1 # 增加当前任务的轮次

            except Exception as e:
                self.log_error(f"任务 {task['task_id']} 调用llm_chat过程中发生错误: {e}")
                task['latest_state'] = "stop"
                # 即使出错也增加 epoch 防止卡死
                task['current_epoch'] += 1

        else:
            # 多任务并行处理
            self.log_info(f"开始并行处理 {len(active_tasks)} 个任务...")
            try:
                # 使用 await 调用异步的 llm_chat_parallel
                responses = await llm_chat_parallel(
                    max_wait_times = max_wait_times,
                    messages_list=messages_list,
                    model_name=model_name,
                    temperature=temperature,
                    ** ({"tools": self.tool_manager.tools} if self.tool_manager.tools else {}),
                    ** ({"tool_choice": tool_choice} if self.tool_manager.tools else {}),
                    ** args
                )

                # 检查返回的 responses 长度是否与 active_tasks 匹配
                if len(responses) != len(active_tasks):
                     self.log_error(f"llm_chat_parallel 返回的响应数量 ({len(responses)}) 与活动任务数量 ({len(active_tasks)}) 不匹配！")
                     # 将所有相关任务标记为停止，避免状态不一致
                     for task in active_tasks:
                          task['latest_state'] = "stop"
                          task['current_epoch'] += 1
                     return # 提前返回

                # 遍历任务和对应的响应
                for task, response in zip(active_tasks, responses):
                     try:
                         # --- 处理每个响应 ---
                         if not response or not hasattr(response, 'choices') or not response.choices:
                              self.log_error(f"任务 {task['task_id']} (并行) 收到无效响应: {response}")
                              task['latest_state'] = "stop" # 标记为停止
                              task['current_epoch'] += 1
                              continue # 处理下一个任务

                         choice = response.choices[0]
                         task['latest_state'] = choice.finish_reason # 更新状态
                         result = choice.message
                         result_dump = result.model_dump() # 这个可能在 result 无效时出错，需要try-except?

                         # 更新内存 - 直接更新 task 引用
                         # task['memory'].add_message(result_dump)

                         # 更新 self.tasks 中的对应项
                         for t in self.tasks:
                              if t['task_id'] == task['task_id']:
                                   t['memory'].add_message(result_dump)
                                   t['latest_state'] = task['latest_state']
                                   t['current_epoch'] = task['current_epoch'] # 同步 epoch
                                   break

                         # 处理工具调用
                         if hasattr(result, 'tool_calls') and result.tool_calls:
                             # _handle_tool_calls 是同步的
                             self._handle_tool_calls(result.tool_calls, task)

                         # 记录 token - 需确认并行接口返回的 response 是否包含 usage
                         # Zhipu 的异步接口返回结果 *不包含* usage 信息！需要单独处理或估算。
                         # if hasattr(response, 'usage') and hasattr(response.usage, 'total_tokens'):
                         #     self.total_tokens += response.usage.total_tokens
                         #     self.log_info(f"任务 {task['task_id']} token 使用: {response.usage.total_tokens} (估算?), 总 token: {self.total_tokens}")
                         # else:
                         #     self.log_warning(f"任务 {task['task_id']} 的并行响应中未找到 token 使用量信息。")

                         task['current_epoch'] += 1 # 增加当前任务的轮次
                     except Exception as inner_e:
                          # 处理单个任务响应解析或处理中的错误
                          self.log_error(f"处理任务 {task['task_id']} 的响应时出错: {inner_e}")
                          task['latest_state'] = "stop"
                          task['current_epoch'] += 1

            except Exception as e:
                # 处理 llm_chat_parallel 本身的调用错误
                self.log_error(f"调用 llm_chat_parallel 过程中发生错误: {e}")
                # 将所有当前活跃的任务标记为停止
                for task in active_tasks:
                    task['latest_state'] = "stop"
                    task['current_epoch'] += 1 # 增加 epoch 防止卡死


    def _handle_tool_calls(self, tool_calls: List[Any], task: Dict[str, Any]):
        """处理工具调用 (保持同步)"""

        # 确保 tool_calls 是列表
        if not isinstance(tool_calls, list):
             self.log_warning(f"任务 {task['task_id']} 的 tool_calls 不是列表: {tool_calls}")
             return

        self.log_debug(f"任务 {task['task_id']} 开始处理 {len(tool_calls)} 个工具调用")
        tool_name = "" # 移到循环外定义，便于错误日志记录
        for index, tool_call in enumerate(tool_calls):
            # 检查 tool_call 结构
            if not hasattr(tool_call, 'function') or not hasattr(tool_call, 'id'):
                 self.log_warning(f"任务 {task['task_id']} 的第 {index+1} 个 tool_call 结构无效: {tool_call}")
                 continue # 跳过无效的工具调用

            try:
                tool = tool_call.function
                tool_call_id = tool_call.id
                tool_name = tool.name
                # 尝试解析 JSON 参数，增加错误处理
                try:
                     tool_args_str = tool.arguments
                     tool_args = json.loads(tool_args_str)
                except json.JSONDecodeError:
                     self.log_error(f"任务 {task['task_id']} 调用工具 '{tool_name}' 的参数 JSON 解析失败: {tool_args_str}")
                     # 可以选择添加一个错误消息给 LLM，或者跳过这个工具调用
                     tool_result = f"Error: Failed to parse arguments JSON for tool {tool_name}."
                     # 添加工具错误消息到内存
                     for t in self.tasks:
                          if t['task_id'] == task['task_id']:
                              t["memory"].add_tool_message(
                                  tool_call_id=tool_call_id,
                                  tool_name=tool_name,
                                  content=tool_result # 返回错误信息
                              )
                              # task['latest_state'] 保持不变，让 LLM 决定下一步
                              break
                     continue # 处理下一个工具调用

                self.log_info(f"任务 {task['task_id']} 调用工具{index+1}: '{tool_name}'，参数: {tool_args}")
                # 调用工具执行
                tool_result = self.tool_manager.call_tool(tool_name, **tool_args)

                # 记录工具结果到对应任务内存
                tool_message_added = False # 标记是否成功添加消息
                for t in self.tasks:
                    if t['task_id'] == task['task_id']:
                        t["memory"].add_tool_message(
                            tool_call_id=tool_call_id,
                            tool_name=tool_name,
                            content=str(tool_result) # 确保结果是字符串
                        )
                        # task['latest_state'] 保持不变，由 LLM 驱动下一步
                        tool_message_added = True
                        break
                if not tool_message_added:
                     self.log_error(f"未能为任务 {task['task_id']} 找到对应的内存来添加工具 {tool_name} 的结果！")


                self.log_info(f"任务 {task['task_id']} 工具'{tool_name}'的返回结果: {str(tool_result)[:200]}...") # 截断过长结果

            except Exception as e:
                # 捕获工具调用执行期间的错误
                self.log_error(f"任务 {task['task_id']} 调用工具 '{tool_name}' 时出错: {e}")
                # 向 LLM 返回错误信息
                error_content = f"Error executing tool '{tool_name}': {str(e)}"
                error_message_added = False
                for t in self.tasks:
                     if t['task_id'] == task['task_id']:
                          t["memory"].add_tool_message(
                               tool_call_id=tool_call_id, # 确保 tool_call_id 已获取
                               tool_name=tool_name,
                               content=error_content
                          )
                          error_message_added = True
                          break
                if not error_message_added:
                     self.log_error(f"未能为任务 {task['task_id']} 找到对应的内存来添加工具 {tool_name} 的错误信息！")


    def is_hit_stop(self, task: Dict[str, Any]) -> bool:
        """判断是否停止某个任务"""
        # finish_reason == 'stop' 表示 LLM 认为完成了当前步骤
        # 但 Agent 的循环是否停止取决于 epoch 和外部条件
        # 这里保持原逻辑：达到最大轮次或者 latest_state 是 'stop'(表示显式停止或错误)
        # Zhipu 返回的 finish_reason 可能有 'tool_calls'，这不应停止 Agent 循环
        current_epoch = task.get('current_epoch', 0)
        latest_state = task.get('latest_state', '')

        if latest_state == "stop": # 显式停止信号或严重错误
             self.log_info(f"任务 {task.get('task_id', 'N/A')} 达到停止状态 '{latest_state}'。")
             return True
        if current_epoch >= self.max_epochs:
             self.log_info(f"任务 {task.get('task_id', 'N/A')} 达到最大轮次 {self.max_epochs}。")
             return True

        # 如果 finish_reason 是 'tool_calls' 或 'length' 等，不应停止
        return False


if __name__ == '__main__':
    class SearchAgent(BaseAgent):
        def __init__(self, max_epochs: int = 5):
            super().__init__(max_epochs=max_epochs)

    # SearchAgent().run_loop(model_name="GLM_4_FLASH", user_input="你叫什么名字",
    #                      system_prompt="你是一个人工智能助手，你叫小明")


    user_input = [
        "你好啊，你是谁",
        "你叫什么名字"
        # "扩散模型的发展现状如何",
    ]

    agent1 = SearchAgent(max_epochs=5)
    responses = agent1.run_loop(model_name="GLM_4_FLASH", user_inputs=user_input,
                         system_prompt="", temperature=0.01)
    print(responses)
